#!/usr/bin/python
# -*-coding:utf-8-*-

import os
import pandas as pd
import numpy as np

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

from dir_info import raw_X_data_dir

def main():
    data_reader = DataReader()
    # data_reader.show_basic_library_db()

    processed_daily_stock_trade_data = data_reader.read_basic_data_table('processed_daily_stock_trade_data')

    processed_daily_stock_trade_data = processed_daily_stock_trade_data.rename(columns={'code': 'stock_code'})

    processed_daily_stock_trade_data = processed_daily_stock_trade_data.set_index(['date', 'stock_code'])

    processed_daily_stock_trade_data['bclose'] = processed_daily_stock_trade_data['close'] * processed_daily_stock_trade_data['back_adjfactor']

    processed_daily_stock_trade_data['swing'] = processed_daily_stock_trade_data['high'] / processed_daily_stock_trade_data['low'] - 1.0

    processed_daily_stock_trade_data['oc_ret'] = processed_daily_stock_trade_data['close'] / processed_daily_stock_trade_data['open'] - 1.0

    processed_daily_stock_trade_data['log_volume'] = np.log(processed_daily_stock_trade_data['volume'])
    processed_daily_stock_trade_data.loc[processed_daily_stock_trade_data['volume']==0, 'log_volume'] = np.nan

    ret = processed_daily_stock_trade_data.groupby('stock_code', as_index=False)['bclose'].pct_change()

    ret.index = ret.index.droplevel(level=0)

    processed_daily_stock_trade_data['ret'] = ret.loc[processed_daily_stock_trade_data.index]

    processed_daily_stock_trade_data = processed_daily_stock_trade_data[
        [
            'ret',
            'oc_ret',
            'swing',
            'log_volume',
            'turnover',
        ]
    ]

    ##
    processed_daily_stock_trade_data = processed_daily_stock_trade_data.astype(np.float32)

    # TODO - 保存
    processed_daily_stock_trade_data.to_hdf(os.path.join(raw_X_data_dir, 'X_daily_price_vol_data.h5'),
                                            key='X_daily_price_vol_data',
                                            mode='w',
                                            format='table')

    print('get daily price&volume data done!\n')



if __name__ == '__main__':
    main()


